import os
import torch
from torch.utils.data import DataLoader, random_split
from .dataset import GenomicDataset
from .model import TMBModel, ChromatinModel
from .train import train_tmb_model, train_model
from .utils import config, logger

def collect_histone_marks(data_dir):
    """收集所有组蛋白修饰的bigWig文件
    
    Args:
        data_dir: 包含组蛋白修饰bigWig文件的目录路径
        
    Returns:
        dict: 组蛋白修饰名称到文件路径的映射
    """
    logger.info(f"从目录 {data_dir} 收集组蛋白修饰文件...")
    histone_marks = {}
    
    if not os.path.exists(data_dir):
        raise ValueError(f"目录不存在: {data_dir}")
    
    # 支持.bw和.bigWig两种后缀
    bigwig_files = [f for f in os.listdir(data_dir) 
                   if f.endswith('.bw') or f.endswith('.bigWig')]
    
    if not bigwig_files:
        raise ValueError(f"在目录 {data_dir} 中没有找到任何bigWig文件")
    
    # 常见的组蛋白修饰名称映射
    mark_name_map = {
        'h3k27ac': 'H3K27ac',
        'h3k4me3': 'H3K4me3',
        'h3k4me1': 'H3K4me1',
        'h3k27me3': 'H3K27me3',
        'h3k9me3': 'H3K9me3',
        'h3k36me3': 'H3K36me3',
        'h3k79me2': 'H3K79me2',
        'ctcf': 'CTCF',
        'polr2a': 'POLR2A',
        'pol2': 'POL2'
    }
    
    for file in bigwig_files:
        # 获取文件名（不含后缀）
        name = os.path.splitext(file)[0].lower()
        
        # 规范化名称
        if name in mark_name_map:
            name = mark_name_map[name]
        else:
            # 如果不在映射表中，将首字母大写
            name = name.upper()
        
        file_path = os.path.join(data_dir, file)
        
        # 检查文件有效性
        if not os.path.exists(file_path):
            logger.warning(f"警告: 文件不存在: {file_path}")
            continue
            
        if os.path.getsize(file_path) == 0:
            logger.warning(f"警告: 文件为空: {file_path}")
            continue
        
        try:
            # 尝试打开文件验证其有效性
            import pyBigWig
            with pyBigWig.open(file_path) as bw:
                if not bw.chroms():
                    logger.warning(f"警告: 文件无效或为空: {file_path}")
                    continue
        except Exception as e:
            logger.error(f"无法打开文件 {file_path}: {str(e)}")
            continue
        
        histone_marks[name] = file_path
        logger.info(f"添加组蛋白修饰: {name} -> {file_path}")
    
    if not histone_marks:
        raise ValueError("没有找到任何有效的组蛋白修饰文件")
    
    logger.info(f"找到 {len(histone_marks)} 个有效的组蛋白修饰文件:")
    for name, path in histone_marks.items():
        logger.info(f"  - {name}: {path}")
    
    return histone_marks

def setup_device():
    """设置计算设备"""
    if torch.cuda.is_available():
        device = torch.device('cuda')
        # 获取GPU信息
        gpu_count = torch.cuda.device_count()
        logger.info(f"\n发现 {gpu_count} 个GPU设备:")
        for i in range(gpu_count):
            gpu_name = torch.cuda.get_device_name(i)
            gpu_memory = torch.cuda.get_device_properties(i).total_memory / 1024**3
            logger.info(f"GPU {i}: {gpu_name} ({gpu_memory:.1f}GB)")
        
        # 设置当前GPU设备
        current_device = torch.cuda.current_device()
        logger.info(f"使用GPU设备 {current_device}: {torch.cuda.get_device_name(current_device)}")
        
        # 清理GPU缓存
        torch.cuda.empty_cache()
    else:
        device = torch.device('cpu')
        logger.warning("\n警告: 未检测到GPU，将使用CPU进行训练（这可能会很慢）")
    
    return device

def run_training(genome_fasta, histone_marks, gene_expression_file=None, window_size=None,
               use_cache=True, output_dir="output",
               dim_reduction_method='none', n_components=50,
               learning_rate=None, epochs=None, batch_size=64,
               target_chromosome=None, target_start=None, target_end=None,
               use_parallel=True):
    """
    运行训练流程
    
    参数:
        genome_fasta: 基因组序列的FASTA文件路径
        histone_marks: 组蛋白修饰标记字典 {标记名称: 文件路径}
        gene_expression_file: 基因表达数据的CSV文件路径（可选）
        window_size: 窗口大小
        use_cache: 是否使用缓存加速数据加载
        output_dir: 输出目录路径
        dim_reduction_method: 降维方法 ('none', 'pca', 'tsne', 'umap')
        n_components: 降维后的维度
        learning_rate: 学习率
        epochs: 训练轮数
        batch_size: 批次大小
        target_chromosome: 目标染色体
        target_start: 目标区域起始位置
        target_end: 目标区域结束位置
        use_parallel: 是否使用GPU-CPU异构并行处理
    
    返回:
        model_path: 保存的模型路径
    """
    try:
        # 设置设备
        device = setup_device()
        
        # 如果histone_marks是目录路径，收集文件
        if isinstance(histone_marks, str):
            histone_marks = collect_histone_marks(histone_marks)
        
        # 检查基因组文件
        if not os.path.exists(genome_fasta):
            raise ValueError(f"基因组文件不存在: {genome_fasta}")
        
        # 创建数据集
        logger.info("\n创建数据集...")
        
        # 记录降维设置
        if dim_reduction_method != 'none':
            logger.info(f"将使用降维方法: {dim_reduction_method}，目标维度: {n_components}")
        else:
            logger.info("不进行降维处理")
        
        dataset = GenomicDataset(
            genome_fasta=genome_fasta,
            histone_marks=histone_marks,
            gene_expression_file=gene_expression_file,
            window_size=window_size,
            use_cache=use_cache,
            cache_dir=os.path.join(output_dir, 'cache'),
            target_chromosome=target_chromosome,
            target_start=target_start,
            target_end=target_end,
            use_parallel=use_parallel
        )
        
        if len(dataset) == 0:
            raise ValueError("数据集为空！请检查输入文件。")
        
        logger.info(f"\n数据集创建完成，共有 {len(dataset)} 个样本")
        
        # 划分训练集和验证集
        logger.info("\n划分训练集和验证集...")
        train_size = int(0.8 * len(dataset))
        val_size = len(dataset) - train_size
        train_dataset, val_dataset = random_split(dataset, [train_size, val_size])
        logger.info(f"训练集: {train_size} 个样本")
        logger.info(f"验证集: {val_size} 个样本")
        
        # 创建数据加载器
        logger.info("\n创建数据加载器...")
        train_loader = DataLoader(
            train_dataset, 
            batch_size=batch_size or config['training_params']['batch_size'], 
            shuffle=True,
            num_workers=0,  # 禁用多进程加载
            pin_memory=False  # 禁用pin_memory
        )
        val_loader = DataLoader(
            val_dataset, 
            batch_size=batch_size or config['training_params']['batch_size'],
            num_workers=0,  # 禁用多进程加载
            pin_memory=False  # 禁用pin_memory
        )
        
        # 创建模型
        logger.info("\n创建模型...")
        # 检查数据集的第一个样本以确定实际特征维度
        first_sample = dataset[0]
        actual_feature_dim = first_sample['features'].shape[0]
        logger.info(f"检测到的特征维度: {actual_feature_dim}")
        
        # 使用实际特征维度而不是硬编码值
        input_dim = actual_feature_dim  # 之前是硬编码为161
        model = TMBModel(
            input_dim=input_dim,
            hidden_dim=config['model_params']['hidden_dim'],
            num_classes=config['model_params']['num_classes'],
            dropout=config['model_params'].get('dropout', 0.1)
        )
        
        # 将模型移动到GPU
        model = model.to(device)
        
        # 如果有多个GPU，使用DataParallel
        if device.type == 'cuda' and torch.cuda.device_count() > 1:
            logger.info(f"使用 {torch.cuda.device_count()} 个GPU进行并行训练")
            model = torch.nn.DataParallel(model)
        
        # 打印模型信息
        total_params = sum(p.numel() for p in model.parameters())
        trainable_params = sum(p.numel() for p in model.parameters() if p.requires_grad)
        logger.info(f"模型总参数量: {total_params:,}")
        logger.info(f"可训练参数量: {trainable_params:,}")
        
        # 训练模型
        logger.info("\n开始训练...")
        trained_model = train_tmb_model(
            model=model,
            train_loader=train_loader,
            val_loader=val_loader,
            num_epochs=epochs or config['training_params']['num_epochs'],
            learning_rate=learning_rate or config['training_params']['learning_rate']
        )
        
        # 保存模型
        os.makedirs(output_dir, exist_ok=True)
        model_save_path = os.path.join(output_dir, 'final_model.pth')
        # 如果使用了DataParallel，需要保存model.module
        if isinstance(trained_model, torch.nn.DataParallel):
            torch.save(trained_model.module.state_dict(), model_save_path)
        else:
            torch.save(trained_model.state_dict(), model_save_path)
        logger.info(f"\n模型训练完成并保存到: {model_save_path}")
        
        # 生成按标签着色的降维可视化图
        if dim_reduction_method != 'none':
            logger.info("\n生成按标签着色的降维可视化图...")
            try:
                # 使用不同的颜色映射生成多张图
                color_maps = [
                    'tab10', 'tab20', 'hsv',      # 分类色彩映射，适合标签
                    'viridis', 'plasma', 'magma', # 连续色谱，视觉效果好
                    'Set1', 'Set2', 'Set3',       # 不同色调的分类映射
                    'Paired', 'Dark2', 'Accent',  # 高对比度的分类映射
                    'Pastel1', 'Pastel2'          # 柔和色调的分类映射
                ]
                for cm in color_maps[:5]:  # 只使用前5个，避免生成太多图片
                    try:
                        dataset.generate_dim_reduction_plot(colormap=cm)
                        logger.info(f"已使用 {cm} 色彩映射生成降维可视化图")
                    except Exception as e:
                        logger.warning(f"使用 {cm} 色彩映射生成图片时出错: {str(e)}")
                logger.info("降维可视化图生成完成，请查看 output/plots 目录")
            except Exception as e:
                logger.warning(f"生成降维可视化图时出错: {str(e)}")
        else:
            logger.info("\n跳过降维可视化图生成 (use_dim_reduction=False)")
        
        # 清理GPU缓存
        if device.type == 'cuda':
            torch.cuda.empty_cache()
        
        return model_save_path
        
    except Exception as e:
        logger.error(f"\n训练过程中出错: {str(e)}")
        # 发生错误时也清理GPU缓存
        if torch.cuda.is_available():
            torch.cuda.empty_cache()
        raise

if __name__ == "__main__":
    main() 